from goby.format import GobyFormat
from lib.cache import file_cache
from lib import get_fields, parse_fields
import requests
from urllib.parse import urljoin
from .client import Client
from .asset import Asset
from .poc import Poc
from .version import version
from .data import ports
import atexit
from config.settings import GOBY_SETTINGS
from export.export import export_to_excel

class Api():
    '''
    goby api 接口
    '''

    def __init__(self, apiurl: dict(type=str, help="url eg http://127.0.0.1") = "http://127.0.0.1:8361", user: str = "",
                 password: dict(type=str, help="密码") = "", **kwargs):
        '''

        :param apiurl: apiurl 接口 defalut 127.0.0.1:8631
        :param user: 用户
        :param password:  密码
        :param kwargs:  requests 的参数可以设置到这里如代理,header 等
        '''
        self._apiurl = apiurl
        self._core = _Core((user, password), **kwargs)
        self.client = Client(self._apiurl, self._core)
        self.asset = Asset(self._apiurl, self._core)
        self.poc = Poc(self._apiurl, self._core)


class _Core():
    def __init__(self, auth: tuple, **kwargs):
        self.kwargs = kwargs
        self._header()
        self.session = requests.session()
        self.session.auth = auth

    def _header(self):
        '''
        处理 header 和 ua 头
        :return:
        '''
        if 'headers' in self.kwargs:
            self.kwargs['headers'].update({
                "Content-Type": "application/json"
            })
            if "User-Agent" not in self.kwargs['headers']:
                self.kwargs['headers'].update({
                    "User-Agent": "goby-api %s" % version
                })
        else:
            self.kwargs['headers'] = {
                "Content-Type": "application/json",
                "User-Agent": "goby-api %s" % version
            }
        if 'verify' not in self.kwargs:
            self.kwargs['verify'] = False

    def get(self, url: dict(type="", help="这是一个url"), data: dict) -> dict:
        '''

        :param url:
        :param data:
        :return:
        '''
        errinfo = {"statusCode": 200, "messages": "", "data": None}
        if data:
            data = "?"+"&".join(["%s=%s" % (k, v)for k, v in data.items()])
        else:
            data = ""
        url = urljoin(url, data)

        try:
            content = self.session.get(url,  **self.kwargs)
            if content.status_code == 200:
                return content.json()
            else:
                raise ValueError("status_code %s is error " %
                                 (content.status_code))
        except Exception as e:
            errinfo["messages"] = "goby-api error info: " + str(e)
            return errinfo

    def get_content(self, url: dict(type="", help="这是一个url"), data: dict) -> str:
        '''

        :param url:
        :param data:
        :return:
        '''
        errinfo = {"statusCode": 200, "messages": "", "data": None}
        if data:
            data = "?"+"&".join(["%s=%s" % (k, v)for k, v in data.items()])
        else:
            data = ""
        url = urljoin(url, data)

        try:
            content = self.session.get(url,  **self.kwargs)
            if content.status_code == 200:
                return content.content
            else:
                raise ValueError("status_code %s is error " %
                                 (content.status_code))
        except Exception as e:
            errinfo["messages"] = "goby-api error info: " + str(e)
            return errinfo

    def post(self, url: dict(type="", help="这是一个url"), data: dict = None, json=None) -> dict:
        '''
        post 请求数据
        :param url: api
        :param data: json数据
        :param kwargs:
        :return:
        '''
        errinfo = {"statusCode": 200, "messages": "", "data": None}
        try:
            content = self.session.post(
                url, data=data, json=json, **self.kwargs)
            if content.status_code == 200:
                return content.json()
            else:
                raise ValueError("status_code %s is error " %
                                 (content.status_code))
        except Exception as e:
            errinfo["messages"] = "goby-api error info: " + str(e)
            return errinfo


class Goby:
    def __init__(self):
        self.goby = Api(GOBY_SETTINGS['api_url'], GOBY_SETTINGS['user'], GOBY_SETTINGS['password'])

        @atexit.register
        def cleanup():
            self.goby._core.session.close()
        
    @property
    def __task_fields(self):
        return [ 
            'taskId',
            'name',
            'targets',
            'state',
            'progress',
            'total.assets',
            'total.vulnerabilities',
            'total.allassets',
            'total.allips',
            'total.allports',
         ]
    
    def check_alive(self) -> bool:
        """存活性检测
        
        Return:
            存活返回True,否则,返回False
        
        """
        alives = self.goby.client.check()
        if alives['statusCode'] == 200:
            return True
        return False
    
    @property
    def tasks(self) -> list:
        tasks = self.goby.client.tasks()
        status = tasks.get('statusCode', 400)
        if status == 200:
            tasks = parse_fields(tasks['data'], self.__task_fields)
            return tasks
        return []
    
    @file_cache()
    def add_targets(self, targets: list, task_name='') -> str:
        """添加待扫描的IP
        
        Args:
            targets: IP组成的列表
        Return:
            对应添加task的id
        """
        tasks = self.goby.client.startScan(targets, ports=ports.Enterprise, taskname=task_name)
        status = tasks['statusCode']
        if status == 200:
            return tasks['data']['taskId']
        return ''
    
    def remove_targets(self, task_ids):
        for task_id in task_ids:
            try:
                removed = self.goby.client.deleteTask(task_id)
                print(removed)
            except:
                print(f'{task_id}删除失败')
                continue
            
    def get_statistic(self, task_id) -> dict:
        statistic = self.goby.client.getStatisticsData(task_id)
        status = statistic.get('statusCode', 400)
        if status == 200:
            return statistic['data']
        return {}
    
    def get_asset(self, task_id):
        asserts = self.goby.asset.assetSearch(f'taskid="{task_id}"')
        return asserts

    def export(self, task_id, save_path):
        results = [{'goby_task': task_id}]
        export_to_excel(results, xlsx_name='goby_test', sheet_name='Goby(扫描中)', save_path=save_path)

    def export_assets(self, task_id, save_path, sheet_name='Goby(asset)'):
        assets = self.get_asset(task_id)
        assets = GobyFormat().assets_format(assets['data'])
        export_to_excel(assets, xlsx_name='goby_asset_test', sheet_name=sheet_name, save_path=save_path)
        
    
    def run(self, ips, save_path):
        task_id = self.add_targets(ips)
        results = [{'goby_task': task_id}]
        export_to_excel(results, xlsx_name='goby_test', sheet_name='goby',save_path=save_path)
        
    
